// memory_consumer.go
package main

import (
	"flag"
	"fmt"
	"math/rand"
	"runtime"
	"strconv"
	"time"
)

// 定义日志函数
func logInfo(message string) {
	logTime := time.Now().Format("2006-01-02 15:04:05")
	fmt.Printf("[%s] info %s\n", logTime, message)
}

// 控制内存占用的 Goroutine
func memoryConsumer(maxMemoryBytes int64) {
	for {
		// 随机选择内存占用
		memory := rand.Int63n(maxMemoryBytes)
		// 分配内存
		memoryBlock := make([]byte, memory)
		for i := range memoryBlock {
			memoryBlock[i] = byte(rand.Intn(256))
		}
		// 打印当前占用的内存
		logInfo(fmt.Sprintf("Using %d KB memory", memory/1024))
		// 随机等待一段时间
		sleepDuration := time.Duration(rand.Intn(5)+3) * time.Minute
		time.Sleep(sleepDuration)
		// 释放内存
		memoryBlock = nil
		// 强制垃圾回收
		runtime.GC()
	}
}

func main() {
	// 定义命令行参数
	var maxMemory string

	flag.StringVar(&maxMemory, "m", "1G", "Maximum memory to use (e.g., 1G, 512M)")

	flag.Parse()

	// 解析最大内存
	var maxMemoryBytes int64
	switch maxMemory[len(maxMemory)-1] {
	case 'G', 'g':
		memValue, err := strconv.Atoi(maxMemory[:len(maxMemory)-1])
		if err != nil {
			logInfo(fmt.Sprintf("Invalid memory format: %v", err))
			return
		}
		maxMemoryBytes = int64(memValue) * 1024 * 1024 * 1024
	case 'M', 'm':
		memValue, err := strconv.Atoi(maxMemory[:len(maxMemory)-1])
		if err != nil {
			logInfo(fmt.Sprintf("Invalid memory format: %v", err))
			return
		}
		maxMemoryBytes = int64(memValue) * 1024 * 1024
	default:
		logInfo("Invalid memory format. Use G or M.")
		return
	}

	logInfo(fmt.Sprintf("Using up to %s memory", maxMemory))

	// 启动内存占用的 Goroutine
	memoryConsumer(maxMemoryBytes)
}